package com.catmiao.ai.controller;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import io.grpc.Metadata;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.grpc.Collections;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;

import static dev.langchain4j.store.embedding.filter.MetadataFilterBuilder.metadataKey;

@RestController
@Slf4j
public class EmbeddingController {

    @Resource
    private EmbeddingStore<TextSegment> embeddingStore;

    @Resource
    private QdrantClient qdrantClient;

    @Resource
    private EmbeddingModel embeddingModel;

    @GetMapping("/embedding/embed")
    public String embed(){

        String prompt = """
                    咏鸡
                鸡鸣破晓光，
                红冠映朝阳。
                金羽披霞彩，
                昂首步高岗。
                """;

        // 文本向量化
        Response<Embedding> embeddingResponse = embeddingModel.embed(prompt);
        System.out.println(embeddingResponse);

        return embeddingResponse.content().toString();
    }

    /**
     * qdrant客户端创建实例
     */
    @GetMapping("/embedding/createCollection")
    public void createCollection(){

        Collections.VectorParams vectorParams = Collections.VectorParams.newBuilder()
                .setDistance(Collections.Distance.Cosine)
                .setSize(1024)
                .build();

        qdrantClient.createCollectionAsync("test-qdarnt-01", vectorParams);
    }


    /**
     * 往向量数据库中存储数据
     */
    @GetMapping("/embedding/setData")
    public void setData(){
        String prompt = """
                    咏鸡
                鸡鸣破晓光，
                红冠映朝阳。
                金羽披霞彩，
                昂首步高岗。
                """;

        TextSegment textSegment = TextSegment.from(prompt);
        textSegment.metadata().put("author","Monster");

        Embedding embeddingContent = embeddingModel.embed(textSegment).content();
        String result = embeddingStore.add(embeddingContent, textSegment);

        System.out.println(result);
    }

    /**
     * 查找
     */
    @GetMapping("/embedding/query1")
    public void query1(){
        Embedding queryEmbedding = embeddingModel.embed("咏鸡说的是什么").content();

        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                .maxResults(1)
                .build();

        EmbeddingSearchResult<TextSegment> result = embeddingStore.search(embeddingSearchRequest);
        System.out.println(result.matches().get(0).embedded().text());
    }

    @GetMapping("/embedding/query2")
    public void query2(){
        Embedding queryEmbedding = embeddingModel.embed("咏鸡").content();

        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                .filter(metadataKey("author").isEqualTo("Monster2"))
                .maxResults(1)
                .build();

        EmbeddingSearchResult<TextSegment> result = embeddingStore.search(embeddingSearchRequest);
        System.out.println(result.matches().get(0).embedded().text());
    }

}
